import numpy as np
import matplotlib.pyplot as plt

#Returns molality of sulfuric acid in mol/kg based on wt% sulfuric acid
def Get_MX(wtX,muX,rhoX,rho0):                            
    MX = wtX/((1.0-wtX)*muX)
    return MX

wtX_mod,z_mod = np.loadtxt('acidity_mod.dat',usecols=(0,1),skiprows=2,unpack=True)  #wt% H2SO4 as a function if atmospheric height, from Dai+2022.

wtX_rhom,z_rhom = np.loadtxt('Rhomb_AFS_Venus.dat',usecols=(3,0),skiprows=1,unpack=True) #rhomboclase stability curve.

#Temperature profile for Venus:
z,T,p = np.loadtxt('pT-Venus.dat',usecols=(0,1,4),skiprows=2,unpack=True) #height in km, temperature in K, pressure in bar

T = T - 273.15              #Convert temperature to celsius

rhoX = 1.83                 #Density of sulfuric acid in kg/L
rho0 = 1.00                 #Density of water in kg/L

muX = 0.098                 #Molar mass of sulfuric acid in kg/mol

MX_mod = Get_MX(wtX_mod,muX,rhoX,rho0)      #Central data in molality, based on centeral data in wt% H2SO4

MX_rhom = Get_MX(wtX_rhom,muX,rhoX,rho0)    #Rhomboclase molality from the stability curve.

MX = np.interp(z,z_mod,MX_mod)

plt.plot(MX,z,'k-')         #Central data molality (mol H2SO4/kg H2O)
plt.plot(MX+60.0,z,'k--')   #Estimated upper limit, based on Dai+2022
plt.plot(MX-29.5,z,'k--')   #Estimated lower limit, based on Dai+2022

#New interpolation of altitude to get the region bounds to stay within the lower bounds of H2SO4 molality.
z_new = np.arange(48.0,62.00001,0.0001) #Altitude (km)

MX_rhom_new = np.interp(z_new,z[::-1],MX_rhom[::-1])        #Rhomboclase molality
MX_new = np.interp(z_new,z[::-1],MX[::-1])                  #lower limit molality

cutoff = 84000                                      #Set to make sure the shaded region terminates at ~62.5 km

plt.fill_betweenx(z_new[cutoff:],MX_rhom_new[cutoff:],MX_new[cutoff:]-29.5,color='y',alpha=0.3)

#Labels for wt% H2SO4
plt.text(4.7,63.0,'35wt%',color='black',fontsize=10)
plt.text(9.0,63.0,'50wt%',color='black',fontsize=10)
plt.text(26.0,63.0,'75wt%',color='black',fontsize=10)
plt.text(80.0,63.0,'90wt%',color='black',fontsize=10)
plt.text(450.0,63.0,'98wt%',color='black',fontsize=10)

#Label for rhomboclase
plt.text(7.5,60.5,'Rhomboclase',color='#8B8000',fontsize=11)

plt.xlabel('molality of H$_2$SO$_4$ (mol H$_2$SO$_4$ / kg H$_2$O)')
plt.ylabel('height in km')

plt.xscale('log')
plt.xlim(3.0,1000.0)
plt.ylim(45.0,65.0)

#Figure file:
plt.savefig('./Molality-Venus.pdf',bbox_inches='tight')

#Output file below, to check plot data against original data:

fout = open('./MolDat_Venus.dat','w')

fout.write('z(km)   p(bar)   T(degC)   min_m_H2SO4(mol/kg)   max_m_H2SO4(mol/kg')

for i in range(len(z)):
    fout.write('\n')
    fout.write('   ')
    fout.write('%.1f' % z[i])
    fout.write('   ')
    fout.write('%.3f' % p[i])
    fout.write('   ')
    fout.write('%.1f' % T[i])
    fout.write('   ')
    fout.write('%.1f' % (MX[i]-29.5))
    fout.write('   ')
    fout.write('%.1f' % (MX[i]+60.0))

fout.close()
